package com.wurd.flink;

import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessAllWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.junit.Test;

import java.time.Duration;
import java.util.*;

public class Demo {

    @Test
    public void test4() throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        List<Tuple2<String, Integer>> data = genData();
        DataStreamSource<Tuple2<String, Integer>> source = env.fromCollection(data);
        WatermarkStrategy<Tuple2<String, Integer>> watermarkStrategy = getWatermarkStrategy();
        source
                .assignTimestampsAndWatermarks(watermarkStrategy)
                .keyBy(0)
//                .sum(1)

                .timeWindowAll(Time.seconds(5))

                .process(new ProcessAllWindowFunction<Tuple2<String, Integer>, Object, TimeWindow>() {
                    @Override
                    public void process(Context context, Iterable<Tuple2<String, Integer>> iterable, Collector<Object> collector) {
                        System.out.println();
                        Tuple2 t = windowDataFormat(iterable);
                        collector.collect(t);
                    }
                })

//                .apply(new AllWindowFunction<Tuple2<String, Integer>, Object, TimeWindow>() {
//                    @Override
//                    public void apply(TimeWindow timeWindow, Iterable<Tuple2<String, Integer>> iterable, Collector<Object> collector) throws Exception {
//                        System.out.println();
//                        Tuple2 t = windowDataFormat(iterable);
//                        collector.collect(t);
//                    }
//                })

                .print();

        System.out.println();
        env.execute("test");
    }

    @Test
    public void test3() throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        List<Tuple2<String, Integer>> data = genData();
        DataStreamSource<Tuple2<String, Integer>> source = env.fromCollection(data);
        WatermarkStrategy<Tuple2<String, Integer>> watermarkStrategy = getWatermarkStrategy();
        source
                .assignTimestampsAndWatermarks(watermarkStrategy)
                .keyBy(0)
//                .sum(1)
                .windowAll(TumblingEventTimeWindows.of(Time.seconds(5)))
//                .windowAll(TumblingProcessingTimeWindows.of(Time.seconds(5)))
                .process(new ProcessAllWindowFunctionA())

                .print();

        System.out.println();
        env.execute("test");
    }

    private WatermarkStrategy<Tuple2<String, Integer>> getWatermarkStrategy() {
        SerializableTimestampAssigner<Tuple2<String, Integer>> timestampAssigner =
                (element, recordTimestamp) -> {
                    try {
                        Thread.sleep(1000);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                    long t = System.currentTimeMillis();
                    return t;
                };
        return WatermarkStrategy.<Tuple2<String, Integer>>forBoundedOutOfOrderness(
                Duration.ofSeconds(10)).withTimestampAssigner(timestampAssigner);
    }

    private List<Tuple2<String, Integer>> genData() {
        List<Tuple2<String, Integer>> data = new ArrayList<>();
        data.add(Tuple2.of("1", 11));
        data.add(Tuple2.of("1", 11));
        data.add(Tuple2.of("2", 22));
        data.add(Tuple2.of("2", 22));
        data.add(Tuple2.of("3", 33));
        data.add(Tuple2.of("3", 33));
        data.add(Tuple2.of("4", 44));
        data.add(Tuple2.of("4", 44));
        data.add(Tuple2.of("5", 55));
        data.add(Tuple2.of("5", 55));
        data.add(Tuple2.of("1", 11));
        data.add(Tuple2.of("1", 11));
        data.add(Tuple2.of("2", 22));
        data.add(Tuple2.of("2", 22));
        data.add(Tuple2.of("3", 33));
        data.add(Tuple2.of("3", 33));
        data.add(Tuple2.of("4", 44));
        data.add(Tuple2.of("4", 44));
        data.add(Tuple2.of("5", 55));
        data.add(Tuple2.of("5", 55));
        return data;
    }

    static class ProcessAllWindowFunctionA extends ProcessAllWindowFunction<Tuple2<String, Integer>, Object, TimeWindow> {
        @Override
        public void process(Context context, Iterable<Tuple2<String, Integer>> iterable, Collector<Object> collector) {
            System.out.println();
            Tuple2 t = windowDataFormat(iterable);
            collector.collect(t);
        }
    }

    private static Tuple2 windowDataFormat(Iterable<Tuple2<String, Integer>> iterable) {
        final Tuple2 t = Tuple2.of("res", 0);
        iterable.forEach(tuple2 -> {
            System.out.println(tuple2);
            t.f1 = (Integer) t.f1 + tuple2.f1;
        });
        return t;
    }

    @Test
    public void test2() throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        new ServerSocketDemo().createSocket();
        DataStreamSource<String> d = env.socketTextStream("localhost", 3333);
        d.print();
        env.execute("socket");
        Thread.currentThread().join();

    }

    @Test
    public void test1() throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        List<Map<String, Integer>> list = new ArrayList<>();
        Map<String, Integer> map = new HashMap<>();
        map.put("one", 1);
        list.add(map);
        list.add(map);
        list.add(map);
        list.add(map);
        list.add(map);
        DataStreamSource<Map<String, Integer>> dataStream = env.fromCollection(list);

//        dataStream.map(new MapFunction<Map<String, Integer>, Object>() {
//            @Override
//            public Object map(Map<String, Integer> map) throws Exception {
//                map.put("one",11);
//                return map;
//            }
//        }).print();
        SingleOutputStreamOperator<Object> d = dataStream.flatMap(new FlatMapFunction<Map<String, Integer>, Object>() {
            @Override
            public void flatMap(Map<String, Integer> map, Collector<Object> collector) throws Exception {
                Thread.sleep(500);
                int r = new Random().nextInt(5);
                System.out.println(r);
                if (r > 2) {
                    map.put("two", 22);
                } else {
                    map.put("two", 33);
                }
                collector.collect(map);
            }
        });
        KeyedStream<Object, Object> dd = d.keyBy(new KeySelector<Object, Object>() {
            @Override
            public Object getKey(Object o) throws Exception {
                Map a = ((Map) o);
                return a.get("two");
            }
        });
        SingleOutputStreamOperator<Object> ddd = dd.reduce(new ReduceFunction<Object>() {
            @Override
            public Object reduce(Object o, Object o1) throws Exception {
                Map a = ((Map) o);
                Map b = ((Map) o1);
                b.put("one", (Integer) a.get("one") + (Integer) b.get("one"));
                return b;
            }
        });
        ddd.print();

        env.execute("test");

    }
}
